import faiss
import numpy as np
import torch
from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import cv2
import time
 
class FaissSearch:
    def __init__(self):
        #Load the model and processor
        self.device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
        self.processor = AutoImageProcessor.from_pretrained('facebook/dinov2-small')
        self.model = AutoModel.from_pretrained('facebook/dinov2-small').to(self.device)
        self.index = faiss.read_index("vector.index")

    def exact_feature_save(self):
        pass

    def search(self,image):

        # Image
        #Extract the features
        time_start = time.time()
        with torch.no_grad():
            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
            outputs = self.model(**inputs)
        print("time_read",time.time()-time_start)
        #Normalize the features before search
        embeddings = outputs.last_hidden_state
        embeddings = embeddings.mean(dim=1)
        vector = embeddings.detach().cpu().numpy()
        vector = np.float32(vector)
        faiss.normalize_L2(vector)
        
        #Read the index file and perform search of top-3 images
        
        d,i = self.index.search(vector,1)
        return d,i

if __name__ == "__main__":
    img = "test.png"
    img2 = "test2.png"
    # image = Image.open(img)
    image = cv2.imread(img)
    image2 = cv2.imread(img2)
    # image = image.resize((60,40))
    faisssearch = FaissSearch()
    faisssearch.search([image,image2])